feat: add oauth_token_cache_enabled to control kernel U2M on-disk token cache - #932
Conversation
…en cache Adds the oauth_token_cache_enabled connect() kwarg, forwarded to the pyo3 Session token_cache_enabled field on the oauth-u2m path. U2M-only; disabled by default (matches Thrift no-persistence posture); enable-flag only; experimental_oauth_persistence is untouched. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
GitHub Actions
Inactive
…TION_PARAMETERS.md The connector forwards bool(None)=False on unset, so omitting the kwarg disables the on-disk cache; it does not inherit the kernel's enabled-by-default. Co-authored-by: Isaac Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
There was a problem hiding this comment.
Pull request overview
Adds a new connection kwarg to control whether the kernel backend persists OAuth U2M refresh tokens to an encrypted on-disk cache, and threads that option through the kernel auth bridge to the kernel Session kwargs.
Changes:
- Introduces
oauth_token_cache_enabled: bool | Noneas asql.connect(...)kwarg and documents it. - Plumbs the option through session/backend wiring so the kernel U2M path forwards
token_cache_enabled. - Expands unit tests to validate defaulting/forwarding behavior and to ensure the flag is not forwarded on PAT/M2M paths.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
tests/unit/test_kernel_auth_bridge.py |
Adds coverage for defaulting and forwarding of token_cache_enabled on U2M only. |
src/databricks/sql/session.py |
Threads oauth_token_cache_enabled through kernel auth option collection (comment currently mismatches behavior). |
src/databricks/sql/client.py |
Documents the new connect kwarg (doc currently mismatches behavior). |
src/databricks/sql/backend/kernel/auth_bridge.py |
Forwards oauth_token_cache_enabled to the kernel’s token_cache_enabled (current coercion can mis-handle non-bool inputs). |
CONNECTION_PARAMETERS.md |
Adds the new parameter to the connection parameter reference table. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
![peco-review-bot[bot]](https://avatars.githubusercontent.com/in/3801177?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: 2 Medium
Solid, well-tested change threading a new U2M-only oauth_token_cache_enabled kwarg to the kernel. Core logic and unit coverage look correct and match the disable-by-default design. Two medium doc/comment mismatches: the client.py docstring and the session.py comment both claim an unset value inherits the kernel's (enabled) default, but the code forwards False (disabled) — the reverse of the actual, and PR-intended, behavior.
eric-wang-1990
added
the
engineer-bot
GitHub Actions
Inactive
Addresses: - #3858601138 at src/databricks/sql/backend/kernel/auth_bridge.py:468 - #3858601153 at src/databricks/sql/session.py:193 - #3858601164 at src/databricks/sql/client.py:234 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Verdict: 1 Medium
Clean, well-scoped change that threads a new oauth_token_cache_enabled kwarg through to the kernel U2M branch, with thorough unit coverage (unset/False/True, truthy/falsey strings via _coerce_bool, and non-forwarding on PAT/M2M). Verdict: looks good — 1 medium, and it's the author's own documented do-not-merge blocker (the unconditional token_cache_enabled forward hard-couples to unreleased kernel PR #283, so merging before the pin bump breaks every U2M connect with a TypeError).
Minor nits (not filed inline): (1) _coerce_bool silently returns False for unknown types, diverging from the fail-loud ProgrammingError pattern in sibling helpers _normalize_scopes/_coerce_redirect_port — defensible as a fail-safe default and explicitly documented, but inconsistent. (2) A couple of the new explanatory comment lines in session.py/auth_bridge.py exceed the repo's 100-char limit (CONTRIBUTING.md). (3) The kernel_auth_kwargs docstrin
[...truncated to keep verdict scannable]
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — clean, well-tested addition of the oauth_token_cache_enabled kwarg. _coerce_bool correctly orders the bool-before-int check (so real booleans aren't misread), handles falsey DSN/env strings like "False", and the flag is correctly scoped to the U2M branch and excluded from PAT/M2M/JWT (all covered by unit tests). One low note on the known forward-compat merge gate. Nit: the client.py docstring and CONNECTION_PARAMETERS.md cite only the Linux cache path (~/.config/databricks-sql-kernel/oauth/); the PR description notes macOS uses ~/Library/Application Support/..., so the docs are slightly incomplete for macOS users.
…bled as Optional[bool]
oauth_token_cache_enabled is a typed connection kwarg like the connector's
other booleans (use_cloud_fetch, _use_arrow_native_complex_types, ...), none
of which coerce string inputs. Replace the one-off _coerce_bool with
`opts.get("oauth_token_cache_enabled") is True`: only a real True enables the
kernel's on-disk U2M token cache, and unset (None)/False forward an explicit
False so the kernel — whose own default is enabled — stays disabled by
default. Any non-bool value fails safe to disabled rather than silently
enabling.
Matches the nodejs connector's `tokenCacheEnabled ?? false` and removes the
inconsistency of coercing this one flag while every sibling boolean is passed
raw. Tests: the string-coercion cases are replaced by one asserting non-bool
inputs never enable the cache.
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
GitHub Actions
Failure
GitHub Actions
Inactive
GitHub Actions
Inactive
…nabled) The token_cache_enabled field this PR forwards to the kernel exists in the pyo3 Session only as of kernel #283. The prior pin (ad78a5b) lacks it, so the flag reached a kernel that ignored it. Bump to 628abd6 (kernel main HEAD): its pyo3 Session signature is a backward-compatible superset of ad78a5b's (identical azure-SP params, plus token_cache_enabled/_passphrase), so the kernel-e2e CI builds a kernel that honors the flag. Verified E2E (U2M): unset => no on-disk cache; True => encrypted cache file written. Published databricks-sql-kernel pin stays ^0.2.0 (no 0.2.1 pyo3 wheel yet); KERNEL_REV is what the kernel-e2e workflow builds the wheel from. Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
Resolves the KERNEL_REV conflict: main pins dd810d6 (#288, request timeout); this branch pins 628abd6 (#283, token cache). dd810d6 is an ancestor of 628abd6, so 628abd6 carries both the request-timeout and token-cache kernel changes — keep 628abd6. All other files auto-merged (disjoint regions: token-cache vs socket-timeout). Signed-off-by: eric-wang-1990 <e.wang@databricks.com>
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
There was a problem hiding this comment.
Verdict: 1 Low
Clean, well-tested addition of the oauth_token_cache_enabled kwarg. The bridge correctly forwards token_cache_enabled only on the oauth-u2m branch as opts.get(...) is True, implementing the documented disable-by-default posture; PAT/M2M/JWT paths are unaffected and covered by tests. No correctness or coverage concerns. One low doc-accuracy note inline. Note the PR is a draft blocked on the databricks-sql-kernel pin bump (already called out in the description) — not re-raised here.
Addresses: - #3860320318 at src/databricks/sql/client.py:230 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
What & why
Adds an
oauth_token_cache_enabled: bool | Noneconnection kwarg to control the kernel backend's U2M on-disk OAuth token cache.On the kernel path with OAuth U2M (
auth_type="oauth-u2m"), the Rust kernel can persist the refresh token to an AES-256 encrypted on-disk cache in the OS config dir (~/Library/Application Support/databricks-sql-kernel/oauth/on macOS,~/.config/databricks-sql-kernel/oauth/on Linux), so a later process skips the interactive browser login. This threads asql.connect(...)kwarg through to the pyo3Session(token_cache_enabled=...)field.Behavior
token_cache_enabled=False— matching Thrift's in-memory (no on-disk persistence) posture, so migrating onto the kernel path doesn't silently start writing tokens to disk. Enabling is opt-in viaoauth_token_cache_enabled=True.experimental_oauth_persistence(customOAuthPersistencestore) is untouched — this new flag is a separate control and the kernel owns its own token lifecycle.Dependency / why this is a draft
Consumes the pyo3 kwarg
token_cache_enabledfrom databricks-sql-kernel PR #283, which is not yet released. Thedatabricks-sql-kernelpin is intentionally left at^0.2.0(a bump to a nonexistent release would break install). Do not merge until:databricks-sql-kernelwheel ships with the kwarg, andNote: because the connector always forwards
token_cache_enabledon U2M, merging before the pin bump would raiseTypeError: unexpected keyword argumentagainst the current binding on every U2M connect — so the pin bump must land atomically with (or before) this reaching a released state.Testing
False,True⇒True,False⇒False; applies to bothdatabricks-oauthandazure-oauth; not forwarded on PAT/M2M (green locally — 24 in the token-cache/U2M subset). These mock the pyo3 Session.token_cache_enabledbranch (databricks-sql-kernel reve72b366, PR Client automatically starts terminated clusters without any indication it has done so #283). One connect per process, isolated to this connection's SHA-256 cache-key file (the shared cache dir was otherwise untouched). All three phases passed:oauth_token_cache_enabled=True, key file removed first): browser flow runs (Opening browser for OAuth authorization…),SELECT 1returns1, and the encrypted key file is written at0o600.Token cache hit— no browser, query returns from the cached/refreshed token.oauth_token_cache_enabled=False, key file removed first): browser flow runs, query returns, and no cache file is written — confirming the disable-by-default posture. Unset behaves identically (both forwardtoken_cache_enabled=Falseat the kernel boundary).Related: databricks-sql-kernel PR #283.
This pull request and its description were written by Isaac.